import numpy as np
with open('../data/jaychou_lyrics.txt', encoding='utf8') as f:
    data = f.read()
print(data)
print(len(data))

corpus_chars = data.replace('\n', ' ').replace('\r', ' ')
corpus_chars = corpus_chars[:20000]
idx_to_char = list(set(corpus_chars))
char_to_idx = dict([(char, i) for i,char in enumerate(idx_to_char)])

vocab_size = len(char_to_idx)

print(vocab_size)

corpus_indices = [char_to_idx[char] for char in corpus_chars]
print(corpus_indices)


import random
from mxnet import nd

def data_iter_random(corpus_indices, batch_size, num_steps, ctx=None):
    num_examples = (len(corpus_indices) - 1) // num_steps
    epoch_size = num_examples // batch_size
    example_indices = list(range(num_examples))

    random.shuffle(example_indices)
    def _data(pos):
        return corpus_indices[pos:pos + num_steps]

    for i in range(epoch_size):
        i = i * batch_size
        batch_indices = example_indices[i: i+batch_size]


        data = nd.array(
            np.array([_data(j * num_steps) for j in batch_indices]), ctx=ctx
        )
        label = nd.array(
            np.array([_data(j * num_steps + 1) for j in batch_indices]), ctx=ctx)
        yield data, label


my_seq = list(range(30))

for data, label in data_iter_random(my_seq, batch_size=2, num_steps=3):
    print('data: ', data,'\nlabel', label)

